import pandas as pd
import aiomysql
from datetime import datetime
import json
import asyncio
import config
from async_db import AsyncMysqlDB
from var import media_crawler_db_var, db_conn_pool_var, source_keyword_var
from tools import utils
from db import DatabaseManager

class DataManager:
    def __init__(self):
        self._lock = asyncio.Lock()
        self._db_manager = None
        
    async def _get_db_manager(self):
        """获取数据库管理器实例"""
        if not self._db_manager:
            self._db_manager = await DatabaseManager.get_instance()
        return self._db_manager
        
    async def ensure_connected(self):
        """确保数据库连接可用"""
        db_manager = await self._get_db_manager()
        return await db_manager.ensure_connected()
        
    async def execute_query(self, query: str, params: list = None) -> list:
        """
        执行SQL查询并返回结果
        :param query: SQL查询语句
        :param params: 查询参数
        :return: 查询结果列表
        """
        try:
            db_manager = await self._get_db_manager()
            await db_manager.ensure_connected()
            async_db = db_manager.get_async_db()
            results = await async_db.query(query, *(params or []))
            return results
        except Exception as e:
            utils.logger.error(f"执行查询出错: {str(e)}")
            return []

    async def close(self):
        """关闭数据库连接"""
        if self._db_manager:
            await self._db_manager.close()
            self._db_manager = None

    async def __aenter__(self):
        """异步上下文管理器入口"""
        await self.ensure_connected()
        return self
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        await self.close()

    async def get_tasks(self, platform=None, limit=100):
        try:
            if platform == "dy":
                # 抖音数据
                query = """
                SELECT 
                    a.id, a.user_id, a.nickname, a.aweme_id as content_id, 
                    a.title, a.desc, a.create_time, a.source_keyword,
                    a.liked_count, a.comment_count
                FROM douyin_aweme a
                WHERE 1=1
                """
                if platform:
                    query += " AND source_keyword IS NOT NULL"
                query += " ORDER BY a.create_time DESC LIMIT %s"
                
            elif platform == "wb":
                # 微博数据
                query = """
                SELECT 
                    n.id, n.user_id, n.nickname, n.note_id as content_id,
                    n.content as title, n.content as desc, n.create_time, n.source_keyword,
                    n.liked_count, n.comments_count as comment_count
                FROM weibo_note n
                WHERE 1=1
                """
                if platform:
                    query += " AND source_keyword IS NOT NULL"
                query += " ORDER BY n.create_time DESC LIMIT %s"
                
            elif platform == "xhs":
                # 小红书数据
                query = """
                SELECT 
                    n.id, n.user_id, n.nickname, n.note_id as content_id,
                    n.title, n.desc, n.time as create_time, n.source_keyword,
                    n.liked_count, n.comment_count
                FROM xhs_note n
                WHERE 1=1
                """
                if platform:
                    query += " AND source_keyword IS NOT NULL"
                query += " ORDER BY n.time DESC LIMIT %s"
                
            else:
                # 所有平台数据
                query = """
                (SELECT 
                    a.id, a.user_id, a.nickname, a.aweme_id as content_id,
                    a.title, a.desc, a.create_time, a.source_keyword,
                    a.liked_count, a.comment_count, '抖音' as platform
                FROM douyin_aweme a
                WHERE source_keyword IS NOT NULL)
                UNION ALL
                (SELECT 
                    n.id, n.user_id, n.nickname, n.note_id,
                    n.content, n.content, n.create_time, n.source_keyword,
                    n.liked_count, n.comments_count, '微博' as platform
                FROM weibo_note n
                WHERE source_keyword IS NOT NULL)
                UNION ALL
                (SELECT 
                    n.id, n.user_id, n.nickname, n.note_id,
                    n.title, n.desc, n.time, n.source_keyword,
                    n.liked_count, n.comment_count, '小红书' as platform
                FROM xhs_note n
                WHERE source_keyword IS NOT NULL)
                ORDER BY create_time DESC
                LIMIT %s
                """
            
            data = await self.execute_query(query, [limit])
            return pd.DataFrame(data)
        except Exception as e:
            utils.logger.error(f"[DataManager.get_tasks] Error getting tasks: {str(e)}")
            return pd.DataFrame()
    
    async def get_comments(self, platform, content_id):
        await self.ensure_connected()
        try:
            if platform == "dy":
                query = """
                SELECT * FROM douyin_aweme_comment 
                WHERE aweme_id = %s
                ORDER BY create_time DESC
                """
            elif platform == "wb":
                query = """
                SELECT * FROM weibo_note_comment 
                WHERE note_id = %s
                ORDER BY create_time DESC
                """
            elif platform == "xhs":
                query = """
                SELECT * FROM xhs_note_comment 
                WHERE note_id = %s
                ORDER BY create_time DESC
                """
            
            data = await self.execute_query(query, [content_id])
            return pd.DataFrame(data)
        except Exception as e:
            print(f"Error getting comments: {str(e)}")
            return pd.DataFrame()
    
    async def get_user_relations(self, platform, user_id):
        await self.ensure_connected()
        try:
            if platform == "dy":
                query = """
                SELECT * FROM dy_relation_user 
                WHERE user_id = %s
                """
            elif platform == "wb":
                query = """
                SELECT * FROM weibo_relation_user 
                WHERE user_id = %s
                """
            else:
                return pd.DataFrame()  # 小红书没有关系数据
            
            data = await self.execute_query(query, [user_id])
            return pd.DataFrame(data)
        except Exception as e:
            print(f"Error getting user relations: {str(e)}")
            return pd.DataFrame()
    
    async def get_creator_info(self, platform, user_id):
        await self.ensure_connected()
        try:
            if platform == "dy":
                query = "SELECT * FROM dy_creator WHERE user_id = %s"
            elif platform == "wb":
                query = "SELECT * FROM weibo_creator WHERE user_id = %s"
            elif platform == "xhs":
                query = "SELECT * FROM xhs_creator WHERE user_id = %s"
            
            data = await self.execute_query(query, [user_id])
            return pd.DataFrame(data)
        except Exception as e:
            print(f"Error getting creator info: {str(e)}")
            return pd.DataFrame()
    
    async def export_to_excel(self, data_type, platform=None):
        try:
            if data_type == "帖子数据":
                df = await self.get_tasks(platform, 1000)
            elif data_type == "评论数据":
                # 获取最新的帖子ID
                posts_df = await self.get_tasks(platform, 1)
                if not posts_df.empty:
                    content_id = posts_df.iloc[0]['content_id']
                    df = await self.get_comments(platform, content_id)
                else:
                    df = pd.DataFrame()
            elif data_type == "用户数据":
                # 获取最新的帖子的用户ID
                posts_df = await self.get_tasks(platform, 1)
                if not posts_df.empty:
                    user_id = posts_df.iloc[0]['user_id']
                    df = await self.get_creator_info(platform, user_id)
                else:
                    df = pd.DataFrame()
            else:
                raise ValueError(f"不支持的数据类型: {data_type}")
            
            if df.empty:
                raise ValueError("没有数据可供导出")
                
            filename = f"{data_type}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
            df.to_excel(filename, index=False)
            return filename
        except Exception as e:
            print(f"Error exporting to excel: {str(e)}")
            raise e 
    
    async def get_unique_user_ids(self, platform=None):
        await self.ensure_connected()
        try:
            keyword = source_keyword_var.get()
            if not keyword:
                print("Warning: No source keyword available")
                return []
                
            if platform == "dy":
                query = """
                SELECT user_id 
                FROM (
                    SELECT DISTINCT user_id, create_time
                    FROM douyin_aweme 
                    WHERE user_id IS NOT NULL 
                    AND source_keyword = %s
                ) t
                ORDER BY create_time DESC 
                LIMIT 100
                """
            elif platform == "wb":
                query = """
                SELECT user_id 
                FROM (
                    SELECT DISTINCT user_id, create_time
                    FROM weibo_note 
                    WHERE user_id IS NOT NULL 
                    AND source_keyword = %s
                ) t
                ORDER BY create_time DESC 
                LIMIT 100
                """
            elif platform == "xhs":
                query = """
                SELECT user_id 
                FROM (
                    SELECT DISTINCT user_id, time as create_time
                    FROM xhs_note 
                    WHERE user_id IS NOT NULL 
                    AND source_keyword = %s
                ) t
                ORDER BY create_time DESC 
                LIMIT 100
                """
            else:
                return []
            
            data = await self.execute_query(query, [keyword])
            return [item['user_id'] for item in data]
        except Exception as e:
            print(f"Error getting unique user ids: {str(e)}")
            return []
    
    async def get_unique_user_ids_and_sec_uids(self, platform=None, keyword=None, limit=None):
        """获取唯一用户ID和昵称列表
        
        Args:
            platform: 平台代码
            keyword: 用户昵称关键词（可选）
            limit: 返回结果数量限制（可选）
            
        Returns:
            List[Tuple[str, str]]: 返回 (user_id, nickname) 元组列表
        """
        try:
            await self.ensure_connected()
            
            if platform != "dy":
                utils.logger.warning(f"[DataManager.get_unique_user_ids_and_sec_uids] 不支持的平台: {platform}")
                return []
                
            # 构建基础查询，使用子查询来处理 DISTINCT 和 ORDER BY
            query = """
                SELECT user_id, nickname
                FROM (
                    SELECT DISTINCT user_id, nickname, last_modify_ts
                    FROM dy_creator
                    WHERE 1=1
            """
            params = []
            
            # 添加昵称模糊查询条件
            if keyword:
                query += " AND nickname LIKE %s"
                params.append(f"%{keyword}%")
                
            query += """
                ) t
                ORDER BY last_modify_ts DESC
            """
            
            # 添加数量限制
            if limit:
                query += " LIMIT %s"
                params.append(limit)
                
            # 执行查询
            data = await self.execute_query(query, params)
            
            # 返回结果列表，如果昵称为空则显示为"未知用户"
            return [(item['user_id'], item['nickname'] or '未知用户') for item in data]
            
        except Exception as e:
            utils.logger.error(f"[DataManager.get_unique_user_ids_and_sec_uids] 获取用户列表失败: {str(e)}")
            return []
    
    def __del__(self):
        """Destructor to ensure resources are cleaned up"""
        if self._db_manager:
            import asyncio
            try:
                loop = asyncio.get_event_loop()
                if not loop.is_closed():
                    loop.run_until_complete(self.close())
            except Exception as e:
                utils.logger.warning(f"[DataManager.__del__] Non-fatal error during cleanup: {str(e)}") 